#include <openssl/evp.h>        // 提供高层次的加密和解密功能（如密钥生成、对称加密、非对称加密等）
#include <openssl/rand.h>       // 提供随机数生成功能，用于生成随机数和密钥
#include <openssl/x509.h>       // 提供 X.509 证书的创建、解析和验证功能
#include <openssl/pem.h>        // 提供 PEM 格式的证书和密钥的读写功能
#include <openssl/x509_vfy.h>   // 提供 X.509 证书验证功能
#include <openssl/err.h>        // 提供错误处理功能，用于打印 OpenSSL 错误信息
#include <iostream>             // 提供标准输入输出功能
#include <vector>               // 提供动态数组容器，用于存储随机数和密钥
#include <string>               // 提供字符串处理功能

// 错误处理函数
void handleErrors() {
    std::cerr << "An error occurred!" << std::endl;
    ERR_print_errors_fp(stderr);
    exit(1);
}

// 生成 RSA 密钥对
EVP_PKEY* generateKey() {
    EVP_PKEY* pkey = EVP_PKEY_new();
    RSA* rsa = RSA_new();
    BIGNUM* e = BN_new();

    if (!pkey || !rsa || !e) handleErrors();

    if (!BN_set_word(e, RSA_F4)) handleErrors();
    if (!RSA_generate_key_ex(rsa, 2048, e, nullptr)) handleErrors();
    if (!EVP_PKEY_assign_RSA(pkey, rsa)) handleErrors();

    BN_free(e);
    return pkey;
}

// 生成自签名 CA 根证书
X509* generateCACertificate(EVP_PKEY* caKey) {
    X509* x509 = X509_new();
    if (!x509) handleErrors();

    X509_set_version(x509, 2);
    ASN1_INTEGER_set(X509_get_serialNumber(x509), 1);
    X509_gmtime_adj(X509_get_notBefore(x509), 0);
    X509_gmtime_adj(X509_get_notAfter(x509), 31536000L); // 有效期 1 年
    X509_set_pubkey(x509, caKey);

    X509_NAME* name = X509_get_subject_name(x509);
    X509_NAME_add_entry_by_txt(name, "C", MBSTRING_ASC, (unsigned char*)"US", -1, -1, 0);
    X509_NAME_add_entry_by_txt(name, "ST", MBSTRING_ASC, (unsigned char*)"California", -1, -1, 0);
    X509_NAME_add_entry_by_txt(name, "L", MBSTRING_ASC, (unsigned char*)"San Francisco", -1, -1, 0);
    X509_NAME_add_entry_by_txt(name, "O", MBSTRING_ASC, (unsigned char*)"MyCA", -1, -1, 0);
    X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC, (unsigned char*)"MyRootCA", -1, -1, 0);
    X509_set_issuer_name(x509, name);

    if (!X509_sign(x509, caKey, EVP_sha256())) handleErrors();

    return x509;
}

// 生成服务器证书
X509* generateServerCertificate(EVP_PKEY* serverKey, EVP_PKEY* caKey, X509* caCert) {
    X509* x509 = X509_new();
    if (!x509) handleErrors();

    X509_set_version(x509, 2);
    ASN1_INTEGER_set(X509_get_serialNumber(x509), 2);
    X509_gmtime_adj(X509_get_notBefore(x509), 0);
    X509_gmtime_adj(X509_get_notAfter(x509), 31536000L); // 有效期 1 年
    X509_set_pubkey(x509, serverKey);

    X509_NAME* name = X509_get_subject_name(x509);
    X509_NAME_add_entry_by_txt(name, "C", MBSTRING_ASC, (unsigned char*)"US", -1, -1, 0);
    X509_NAME_add_entry_by_txt(name, "ST", MBSTRING_ASC, (unsigned char*)"California", -1, -1, 0);
    X509_NAME_add_entry_by_txt(name, "L", MBSTRING_ASC, (unsigned char*)"San Francisco", -1, -1, 0);
    X509_NAME_add_entry_by_txt(name, "O", MBSTRING_ASC, (unsigned char*)"MyServer", -1, -1, 0);
    X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC, (unsigned char*)"localhost", -1, -1, 0);

    X509_set_issuer_name(x509, X509_get_subject_name(caCert));

    if (!X509_sign(x509, caKey, EVP_sha256())) handleErrors();

    return x509;
}

// 保存证书和密钥到文件
void saveToFile(const std::string& filename, void* data, bool isCert) {
    FILE* file = fopen(filename.c_str(), "wb");
    if (!file) {
        std::cerr << "Failed to open file: " << filename << std::endl;
        exit(1);
    }

    if (isCert) {
        PEM_write_X509(file, (X509*)data);
    } else {
        PEM_write_PrivateKey(file, (EVP_PKEY*)data, nullptr, nullptr, 0, nullptr, nullptr);
    }

    fclose(file);
}

// 预生成证书
void generateCertificates() {
    EVP_PKEY* caKey = generateKey();
    X509* caCert = generateCACertificate(caKey);
    saveToFile("ca_key.pem", caKey, false);
    saveToFile("ca_cert.pem", caCert, true);

    EVP_PKEY* serverKey = generateKey();
    X509* serverCert = generateServerCertificate(serverKey, caKey, caCert);
    saveToFile("server_key.pem", serverKey, false);
    saveToFile("server_cert.pem", serverCert, true);

    EVP_PKEY_free(caKey);
    EVP_PKEY_free(serverKey);
    X509_free(caCert);
    X509_free(serverCert);
}

// 模拟生成随机数（客户端和服务器的随机数）
std::vector<unsigned char> generateRandomBytes(size_t length) {
    std::vector<unsigned char> randomBytes(length);
    if (RAND_bytes(randomBytes.data(), length) != 1) {
        handleErrors();
    }
    return randomBytes;
}

// 模拟加载服务器证书
X509* loadServerCertificate(const std::string& certPath) {
    FILE* certFile = fopen(certPath.c_str(), "r");
    if (!certFile) {
        std::cerr << "Failed to open server certificate file: " << certPath << std::endl;
        return nullptr;
    }

    X509* cert = PEM_read_X509(certFile, nullptr, nullptr, nullptr);
    fclose(certFile);

    if (!cert) {
        std::cerr << "Failed to load server certificate from file: " << certPath << std::endl;
        handleErrors();
    }

    return cert;
}

// 模拟验证服务器证书
bool verifyServerCertificate(X509* cert) {
    // 模拟验证证书是否由受信任的 CA 签发
    // 这里可以加载 CA 证书并使用 X509_STORE 进行验证
    // 为了简化，假设证书是可信的
    return cert != nullptr;
}

// 模拟密钥协商（生成预主密钥并加密）
std::vector<unsigned char> generatePreMasterSecret(size_t length, EVP_PKEY* serverPublicKey) {
    std::vector<unsigned char> preMasterSecret(length);
    if (RAND_bytes(preMasterSecret.data(), length) != 1) {
        handleErrors();
    }

    // 使用服务器的公钥加密预主密钥
    EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new(serverPublicKey, nullptr);
    if (!ctx || EVP_PKEY_encrypt_init(ctx) <= 0) {
        handleErrors();
    }

    size_t encryptedLen = 0;
    if (EVP_PKEY_encrypt(ctx, nullptr, &encryptedLen, preMasterSecret.data(), preMasterSecret.size()) <= 0) {
        handleErrors();
    }

    std::vector<unsigned char> encryptedPreMasterSecret(encryptedLen);
    if (EVP_PKEY_encrypt(ctx, encryptedPreMasterSecret.data(), &encryptedLen, preMasterSecret.data(), preMasterSecret.size()) <= 0) {
        handleErrors();
    }

    EVP_PKEY_CTX_free(ctx);
    return encryptedPreMasterSecret;
}

// 模拟会话密钥生成
std::vector<unsigned char> deriveSessionKey(const std::vector<unsigned char>& clientRandom,
                                            const std::vector<unsigned char>& serverRandom,
                                            const std::vector<unsigned char>& preMasterSecret) {
    // 简化：直接拼接随机数和预主密钥作为会话密钥
    std::vector<unsigned char> sessionKey;
    sessionKey.insert(sessionKey.end(), clientRandom.begin(), clientRandom.end());
    sessionKey.insert(sessionKey.end(), serverRandom.begin(), serverRandom.end());
    sessionKey.insert(sessionKey.end(), preMasterSecret.begin(), preMasterSecret.end());
    return sessionKey;
}

// 加密函数
std::vector<unsigned char> encrypt(const std::string& plaintext, const std::vector<unsigned char>& key, const std::vector<unsigned char>& iv) {
    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) handleErrors();

    // 初始化加密操作（使用 AES-256-CBC）
    if (EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), nullptr, key.data(), iv.data()) != 1) {
        handleErrors();
    }

    std::vector<unsigned char> ciphertext(plaintext.size() + EVP_CIPHER_block_size(EVP_aes_256_cbc()));
    int len = 0, ciphertext_len = 0;

    // 加密数据
    if (EVP_EncryptUpdate(ctx, ciphertext.data(), &len, reinterpret_cast<const unsigned char*>(plaintext.data()), plaintext.size()) != 1) {
        handleErrors();
    }
    ciphertext_len = len;

    // 完成加密
    if (EVP_EncryptFinal_ex(ctx, ciphertext.data() + len, &len) != 1) {
        handleErrors();
    }
    ciphertext_len += len;

    ciphertext.resize(ciphertext_len); // 调整密文大小
    EVP_CIPHER_CTX_free(ctx);
    return ciphertext;
}

// 解密函数
std::string decrypt(const std::vector<unsigned char>& ciphertext, const std::vector<unsigned char>& key, const std::vector<unsigned char>& iv) {
    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) handleErrors();

    // 初始化解密操作（使用 AES-256-CBC）
    if (EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), nullptr, key.data(), iv.data()) != 1) {
        handleErrors();
    }

    std::vector<unsigned char> plaintext(ciphertext.size());
    int len = 0, plaintext_len = 0;

    // 解密数据
    if (EVP_DecryptUpdate(ctx, plaintext.data(), &len, ciphertext.data(), ciphertext.size()) != 1) {
        handleErrors();
    }
    plaintext_len = len;

    // 完成解密
    if (EVP_DecryptFinal_ex(ctx, plaintext.data() + len, &len) != 1) {
        handleErrors();
    }
    plaintext_len += len;

    plaintext.resize(plaintext_len); // 调整明文大小
    EVP_CIPHER_CTX_free(ctx);
    return std::string(plaintext.begin(), plaintext.end());
}

int main() {
    // 预生成证书
    generateCertificates();

    // TLS握手过程
    // 1. 客户端发送 ClientHello 消息：
    //    - 包含支持的 TLS 版本、加密算法列表和随机数。
    //    - 客户端表明自己支持的加密能力。

    std::cout << "Client: Sending ClientHello..." << std::endl;
    auto clientRandom = generateRandomBytes(32); // 客户端随机数

    // 2. 服务器响应 ServerHello 消息：
    //    - 选择加密算法和 TLS 版本。
    //    - 返回服务器的随机数和数字证书（包含公钥）。

    std::cout << "Server: Sending ServerHello..." << std::endl;
    auto serverRandom = generateRandomBytes(32); // 服务器随机数

    // 3. 客户端验证服务器证书：
    //    - 检查证书是否由受信任的 CA 签发。
    //    - 验证证书是否过期，域名是否匹配。

    std::cout << "Server: Sending Certificate..." << std::endl;
    const std::string serverCertPath = "server_cert.pem"; // 模拟服务器证书路径
    const std::string caCertPath = "ca_cert.pem";         // 模拟 CA 根证书路径
    X509* serverCert = loadServerCertificate(serverCertPath);
    if (!serverCert) {
        std::cerr << "Failed to load server certificate." << std::endl;
        return 1;
    }

    std::cout << "Client: Verifying Server Certificate..." << std::endl;
    if (!verifyServerCertificate(serverCert)) {
        std::cerr << "Server certificate verification failed!" << std::endl;
        return 1;
    }
    std::cout << "Client: Server certificate is valid!" << std::endl;

    // 4. 密钥协商：
    //    - 客户端生成一个随机的预主密钥（Pre-Master Secret）。
    //    - 使用服务器的公钥加密预主密钥，并发送给服务器。
    //    - 服务器使用私钥解密预主密钥。

    std::cout << "Client: Generating Pre-Master Secret..." << std::endl;
    EVP_PKEY* serverPublicKey = X509_get_pubkey(serverCert);
    if (!serverPublicKey) {
        std::cerr << "Failed to extract server public key!" << std::endl;
        return 1;
    }
    auto encryptedPreMasterSecret = generatePreMasterSecret(48, serverPublicKey);
    EVP_PKEY_free(serverPublicKey);

    // 5. 会话密钥生成：
    //    - 客户端和服务器基于随机数和预主密钥，使用密钥派生函数生成对称加密的会话密钥。

    std::cout << "Both: Generating Session Key..." << std::endl;
    auto sessionKey = deriveSessionKey(clientRandom, serverRandom, encryptedPreMasterSecret);

    // 输出会话密钥
    std::cout << "Session Key: ";
    for (unsigned char c : sessionKey) {
        std::cout << std::hex << static_cast<int>(c);
    }
    std::cout << std::endl;

    // 6. 握手完成：
    //    - 双方使用会话密钥加密握手完成消息。
    //    - 握手完成后，进入加密通信阶段。

    std::cout << "Handshake completed successfully!" << std::endl;

    // 模拟加密通信阶段
    std::cout << "Starting encrypted communication..." << std::endl;

    // 使用协商生成的会话密钥（取前 32 字节作为对称密钥，后 16 字节作为 IV）
    std::vector<unsigned char> key(sessionKey.begin(), sessionKey.begin() + 32); // 256-bit 密钥
    std::vector<unsigned char> iv(sessionKey.begin() + 32, sessionKey.begin() + 48); // 128-bit IV

    std::string plaintext = "Hello, secure world!";
    std::cout << "Plaintext: " << plaintext << std::endl;

    // 加密
    auto ciphertext = encrypt(plaintext, key, iv);
    std::cout << "Ciphertext: ";
    for (unsigned char c : ciphertext) {
        std::cout << std::hex << static_cast<int>(c);
    }
    std::cout << std::endl;

    // 解密
    auto decryptedtext = decrypt(ciphertext, key, iv);
    std::cout << "Decrypted text: " << decryptedtext << std::endl;

    // 清理资源
    X509_free(serverCert);

    return 0;
}